import pandas as pd
import plotly.express as px
import plotly.graph_objects as go


def get_employment_num_dic():
    path = f".\劳动力数据源\第五次经济普查法人单位分行业从业人数.xlsx"
    dic = {}
    df = pd.read_excel(path, index_col=0)
    for k,v in df.iterrows():
        if not pd.isna(v.to_list()[-1]):
            dic[k] = int(v.to_list()[-1])
    return dic

def get_category(sheet_name='腾景70行业'):
    path = f".\劳动力数据源\第五次经济普查行业对应70行业40行业.xlsx"
    df = pd.read_excel(path, sheet_name=sheet_name, index_col=0)
    dic = {}
    for k,v in df.iterrows():
        dic[k] = v.dropna().to_list()
    return dic

def get_70_employment_num():
    employment_num_dic = get_employment_num_dic()
    category_70_dic = get_category()
    data_dic = {}
    for k,v in category_70_dic.items():
        data_dic[k] = sum([employment_num_dic[i] for i in v])
    se = pd.Series(data_dic, name='腾景70行业法人单位就业数量')
    return se

def get_40_employment_num():
    employment_num_dic = get_employment_num_dic()
    category_40_dic = get_category(sheet_name='40行业')
    data_dic = {}
    for k,v in category_40_dic.items():
        data_dic[k] = sum([employment_num_dic[i] for i in v])
    se = pd.Series(data_dic, name='40行业法人单位就业数量')
    return se

def get_employment_num(category_dic, name):
    employment_num_dic = get_employment_num_dic()
    data_dic = {}
    for k,v in category_dic.items():
        data_dic[k] = sum([employment_num_dic[i] for i in v])
    se = pd.Series(data_dic, name=name)
    return se

se70 = get_70_employment_num()
se_ratio = se70 / se70.sum()

se40 = get_40_employment_num()
se40_ratio = se40 / se40.sum()

se20 = get_employment_num(get_category(sheet_name='20门类'), '20门类法人单位就业数量')
se20_ratio = se20 / se20.sum()


def tree_plot(se):

    df = pd.DataFrame(se)
    df['country'] = se.index.to_list()
    df['value'] = se.to_list()
    df['parent'] = ['' for i in se.index.to_list()]

    # 创建矩形树图
    fig = go.Figure(go.Treemap(
        labels=df['country'],
        parents=df['parent'],
        values=df['value'],
        marker=dict(
            colors=['lightblue', 'purple', 'red', 'lightpurple', 'lightred', 'lightviolet', 'lightpink', 'lightcoral', 'lightpurple', 'lightcoral', 'lightblue', 'lightgreen', 'lightcyan', 'lightblue', 'lightyellow', 'lightgreen', 'lightblue']  # 可以根据需要自定义颜色列表
        )
    ))

    # 更新布局
    fig.update_layout(
        title='分行业法人单位从业人员',
        margin=dict(t=50, l=25, r=25, b=25)  # 调整边距
    )

    # 显示图形
    return fig

fig = tree_plot(se70)

